from sentence_transformers import SentenceTransformer
import chromadb
from llm.doubao import doubao_qa
from llm.local import ollama_qa
import json

model = SentenceTransformer("all-MiniLM-L6-v2")

client = chromadb.PersistentClient(path="./chromadb_data")
collection = client.get_or_create_collection("rag")


def get_query_embedding(query):
    return model.encode(query).tolist()


def retrieve_related_chunks(query_embedding, n_results=3):
    results = collection.query(query_embeddings=[query_embedding], n_results=n_results)
    related_chunks = results.get("documents")
    if not related_chunks or not related_chunks[0]:
        exit(1)
    return related_chunks[0]


def decompose_complex_query(query):
    decomposition_prompt = f"""
    任务：将用户的复杂问题分解为一系列相关的小问题，以帮助更好地检索相关信息。

    要求：
    - 子问题必须与原始问题高度相关
    - 每个子问题应该明确具体
    - 子问题组合应覆盖原始问题的所有方面

    用户问题：{query}

    请将上述问题分解为3-5个子问题，以JSON格式返回：
    {{
    "sub_questions": [
        "子问题1",
        "子问题2",
        "子问题3"
    ]
    }}
    """
    decomposition_result = ollama_qa(decomposition_prompt)
    jsons_tart = decomposition_result.find("{")
    jsons_end = decomposition_result.rfind("}") + 1
    json_str = decomposition_result[jsons_tart:jsons_end]
    parsed_result = json.loads(json_str)
    sub_questions = parsed_result.get("sub_questions", [])
    return sub_questions


def parallel_retrieve_sub_questions(sub_questions, n_results_per_question=2):
    all_related_chunks = []
    for i, sub_question in enumerate(sub_questions):
        sub_question_embedding = get_query_embedding(sub_question)
        related_chunks = retrieve_related_chunks(
            sub_question_embedding, n_results_per_question
        )
        all_related_chunks.extend(related_chunks)
    unique_chunks = list(set(all_related_chunks))
    return unique_chunks


def integrate_answers(original_query, all_related_chunks):
    context = "\n".join(all_related_chunks)
    integration_prompt = f"""
基于以下检索到的相关信息，请回答用户的原始问题

检索到的相关信息为:
{context}

用户的原始问题为:
{original_query}

请提供一个全面准确的回答，确保
1. 回答覆盖原问题的所有的方面，不要遗漏
2. 基于检索到的信息进行回答
3. 如果信息不足，则明确指出
4. 保持逻辑清晰，结构合理
"""
    final_answer = ollama_qa(integration_prompt)
    return final_answer


if __name__ == "__main__":
    query = input("请输入你的问题:")
    # 智能分解复杂问题
    sub_questions = decompose_complex_query(query)
    # 并行检索子问题
    all_related_chunks = parallel_retrieve_sub_questions(
        sub_questions, n_results_per_question=2
    )
    # 整合答案
    final_answer = integrate_answers(query, all_related_chunks)
    print(f"answer:{final_answer}")
